package com.example.vpn_demo.packet

import android.util.Log
import com.example.vpn_demo.model.HttpRequest
import java.nio.ByteBuffer
import java.io.FileInputStream
import java.io.FileOutputStream
import java.net.Socket

class PacketProcessor(
    private val vpnInput: FileInputStream,
    private val vpnOutput: FileOutputStream,
    private val proxyPort: Int
) {
    companion object {
        private const val TAG = "PacketProcessor"
    }

    fun processPacket(packet: Packet) {
        try {
            when (packet.protocol) {
                Protocol.TCP -> {
                    // 对于 HTTP/HTTPS 流量，转发到本地代理
                    if (packet.destinationPort == 80 || packet.destinationPort == 443) {
                        forwardToProxy(packet)
                    } else {
                        // 非 HTTP/HTTPS 流量直接转发
                        vpnOutput.write(packet.data)
                        vpnOutput.flush()
                    }
                }
                Protocol.UDP -> {
                    // 只记录 DNS 查询
                    if (packet.destinationPort == 53) {
                        Log.i(TAG, "DNS 查询: ${packet.destinationAddress}")
                    }
                    vpnOutput.write(packet.data)
                    vpnOutput.flush()
                }
                else -> {
                    vpnOutput.write(packet.data)
                    vpnOutput.flush()
                }
            }
        } catch (e: Exception) {
            Log.e(TAG, "处理数据包失败", e)
        }
    }

    private fun forwardToProxy(packet: Packet) {
        // 修改目标地址为本地代理
        val modifiedPacket = rewriteDestination(
            packet.data,
            "127.0.0.1",
            proxyPort
        )
        vpnOutput.write(modifiedPacket)
        vpnOutput.flush()
    }

    fun readPacket(): Packet? {
        try {
            val buffer = ByteArray(32767)
            val length = vpnInput.read(buffer)
            if (length <= 0) return null

            // 解析 IP 头
            val ipVersion = (buffer[0].toInt() and 0xF0) shr 4
            val ipHeaderLength = (buffer[0].toInt() and 0x0F) * 4
            val protocol = buffer[9].toInt() and 0xFF

            // 获取源地址和目标地址
            val sourceAddress = getIpAddress(buffer.copyOfRange(12, 16))
            val destinationAddress = getIpAddress(buffer.copyOfRange(16, 20))

            // 如果是 TCP 或 UDP
            if (protocol == 6 || protocol == 17) {
                val transportHeaderOffset = ipHeaderLength
                val sourcePort = getPort(buffer, transportHeaderOffset)
                val destinationPort = getPort(buffer, transportHeaderOffset + 2)

                return Packet(
                    protocol = if (protocol == 6) Protocol.TCP else Protocol.UDP,
                    sourceAddress = sourceAddress,
                    sourcePort = sourcePort,
                    destinationAddress = destinationAddress,
                    destinationPort = destinationPort,
                    data = buffer.copyOfRange(0, length)
                )
            }

            return null
        } catch (e: Exception) {
            Log.e(TAG, "读取数据包失败", e)
            return null
        }
    }

    private fun getIpAddress(bytes: ByteArray): String {
        return bytes.joinToString(".") { (it.toInt() and 0xFF).toString() }
    }

    private fun getPort(buffer: ByteArray, offset: Int): Int {
        return ((buffer[offset].toInt() and 0xFF) shl 8) or
                (buffer[offset + 1].toInt() and 0xFF)
    }

    private fun rewriteDestination(packet: ByteArray, newAddress: String, newPort: Int): ByteArray {
        val ipHeaderLength = (packet[0].toInt() and 0x0F) * 4
        val newPacket = packet.copyOf()
        
        // 修改目标 IP 地址
        val addressBytes = newAddress.split(".").map { it.toInt().toByte() }.toByteArray()
        System.arraycopy(addressBytes, 0, newPacket, 16, 4)
        
        // 修改目标端口
        newPacket[ipHeaderLength + 2] = ((newPort shr 8) and 0xFF).toByte()
        newPacket[ipHeaderLength + 3] = (newPort and 0xFF).toByte()
        
        // 重新计算校验和
        calculateIpChecksum(newPacket)
        calculateTcpChecksum(newPacket)
        
        return newPacket
    }

    private fun calculateIpChecksum(packet: ByteArray) {
        val headerLength = (packet[0].toInt() and 0x0F) * 4
        var sum = 0
        
        // 清除现有校验和
        packet[10] = 0
        packet[11] = 0
        
        // 计算校验和
        var i = 0
        while (i < headerLength) {
            sum += ((packet[i].toInt() and 0xFF) shl 8) or (packet[i + 1].toInt() and 0xFF)
            i += 2
        }
        
        // 处理进位
        while (sum shr 16 > 0) {
            sum = (sum and 0xFFFF) + (sum shr 16)
        }
        
        // 取反
        val checksum = sum.inv() and 0xFFFF
        
        // 写入校验和
        packet[10] = (checksum shr 8).toByte()
        packet[11] = checksum.toByte()
    }

    private fun calculateTcpChecksum(packet: ByteArray) {
        val ipHeaderLength = (packet[0].toInt() and 0x0F) * 4
        val tcpLength = packet.size - ipHeaderLength
        var sum = 0
        
        // 清除现有校验和
        packet[ipHeaderLength + 16] = 0
        packet[ipHeaderLength + 17] = 0
        
        // 计算伪头部校验和
        var i = 12
        while (i < 20) { // 源地址和目标地址
            sum += ((packet[i].toInt() and 0xFF) shl 8) or (packet[i + 1].toInt() and 0xFF)
            i += 2
        }
        sum += 6 // TCP 协议号
        sum += tcpLength
        
        // 计算 TCP 头和数据的校验和
        i = ipHeaderLength
        while (i < packet.size - 1) {
            sum += ((packet[i].toInt() and 0xFF) shl 8) or (packet[i + 1].toInt() and 0xFF)
            i += 2
        }
        if (i < packet.size) { // 处理奇数长度
            sum += (packet[i].toInt() and 0xFF) shl 8
        }
        
        // 处理进位
        while (sum shr 16 > 0) {
            sum = (sum and 0xFFFF) + (sum shr 16)
        }
        
        // 取反
        val checksum = sum.inv() and 0xFFFF
        
        // 写入校验和
        packet[ipHeaderLength + 16] = (checksum shr 8).toByte()
        packet[ipHeaderLength + 17] = checksum.toByte()
    }

    fun writePacket(data: ByteArray) {
        try {
            synchronized(vpnOutput) {
                vpnOutput.write(data)
                vpnOutput.flush()
            }
        } catch (e: Exception) {
            Log.e(TAG, "写入数据包失败", e)
            throw e
        }
    }
} 